"""
LightGBM模型运行脚本
Author: Model Team
Date: 2025-11-16
Description: LightGBM Model for house price prediction
"""

import os
import sys
import json
import pandas as pd
import numpy as np
try:
    import lightgbm as lgb
except ImportError:
    print("[ERROR] LightGBM not installed. Please install with: pip install lightgbm")
    sys.exit(1)

from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import joblib
import warnings
from datetime import datetime

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from model.validation_report_utils import generate_validation_outputs

warnings.filterwarnings('ignore')

class LightGBMModelRunner:
    def __init__(self, config_path='model_config.json'):
        """初始化LightGBM模型运行器"""
        # 确保我们在正确的目录中运行
        script_dir = os.path.dirname(os.path.abspath(__file__))
        os.chdir(script_dir)

        self.config = self.load_config(config_path)
        self.model = None
        self.X_train = None
        self.y_train = None
        self.X_test = None
        self.test_ids = None
        self.feature_names = None
        self.feature_importances = None
        self.metrics = {}

        # 设置输出目录
        self.setup_output_directories()

    def setup_output_directories(self):
        """设置输出目录结构"""
        self.output_dirs = {
            'models': 'outputs/models',
            'submissions': 'outputs/submissions',
            'reports': 'outputs/reports',
            'logs': 'outputs/logs'
        }

        # 创建所有输出目录
        for dir_name, dir_path in self.output_dirs.items():
            os.makedirs(dir_path, exist_ok=True)
            print(f"[OK] 创建目录: {dir_path}")

    def load_config(self, config_path):
        """加载模型配置"""
        try:
            if not os.path.exists(config_path):
                print(f"[ERROR] 配置文件不存在: {config_path}")
                print(f"当前目录中的文件: {os.listdir('.')}")
                sys.exit(1)

            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)

            print(f"[OK] 配置文件加载成功")
            print(f"模型名称: {config['model_info']['model_name']}")
            return config

        except Exception as e:
            print(f"[ERROR] 配置文件加载失败: {e}")
            sys.exit(1)

    def load_data(self):
        """加载数据"""
        print("[INFO] 加载数据...")

        try:
            # 获取项目根目录
            current_dir = os.getcwd()
            script_dir = os.path.dirname(os.path.abspath(__file__))
            model_dir = os.path.dirname(script_dir)
            project_root = os.path.dirname(model_dir)

            # 加载数据
            train_path = os.path.join(project_root, 'feature/train_feature.csv')
            test_path = os.path.join(project_root, 'feature/test_feature.csv')

            if not os.path.exists(train_path) or not os.path.exists(test_path):
                print(f"[ERROR] 数据文件不存在: {train_path} 或 {test_path}")
                return False

            train_data = pd.read_csv(train_path)
            test_data = pd.read_csv(test_path)

            # 准备训练数据
            self.X_train = train_data.drop(['SalePrice', 'Id'], axis=1, errors='ignore')
            self.y_train = np.log1p(train_data['SalePrice'])  # 对数变换
            self.X_test = test_data.drop(['Id'], axis=1, errors='ignore')
            self.test_ids = test_data['Id'] if 'Id' in test_data.columns else None

            # 特征选择
            self.apply_feature_selection()

            self.feature_names = list(self.X_train.columns)

            print(f"[OK] 数据加载成功")
            print(f"   训练集形状: {self.X_train.shape}")
            print(f"   测试集形状: {self.X_test.shape}")
            print(f"   特征数量: {len(self.feature_names)}")
            return True

        except Exception as e:
            print(f"[ERROR] 数据加载失败: {e}")
            return False

    def apply_feature_selection(self):
        """应用特征选择"""
        feature_params = self.config.get('feature_params', {})

        if not feature_params.get('feature_selection', False):
            print("[INFO] 特征选择已禁用，使用所有特征")
            return

        if feature_params.get('use_feature_importance', False):
            self.select_features_by_importance()
        else:
            print("[INFO] 使用默认特征选择方式")

    def select_features_by_importance(self):
        """根据特征重要性选择特征"""
        feature_params = self.config['feature_params']
        importance_file = feature_params.get('feature_importance_file', '../../feature/feature_importance.csv')
        top_n = feature_params.get('top_n_features', 20)

        print(f"[INFO] 使用特征重要性文件: {importance_file}")
        print(f"[INFO] 选择前 {top_n} 个重要特征")

        try:
            # 读取特征重要性文件
            importance_df = pd.read_csv(importance_file)

            # 获取前N个重要特征
            top_features = importance_df.head(top_n)['feature'].tolist()
            print(f"[INFO] 前 {top_n} 个重要特征: {top_features}")

            # 检查特征在数据中是否存在
            available_features = []
            missing_features = []

            for feature in top_features:
                if feature in self.X_train.columns:
                    available_features.append(feature)
                else:
                    missing_features.append(feature)

            if missing_features:
                print(f"[WARNING] 以下特征不存在于数据中: {missing_features}")

            # 选择存在的特征
            if available_features:
                self.X_train = self.X_train[available_features]
                self.X_test = self.X_test[available_features]
                print(f"[OK] 成功选择 {len(available_features)} 个特征")
            else:
                print("[ERROR] 没有找到任何匹配的特征，使用所有特征")

        except Exception as e:
            print(f"[ERROR] 特征重要性选择失败: {e}")
            print("[INFO] 使用所有特征继续训练")

    def train_model(self):
        """训练模型"""
        print("[INFO] 开始训练LightGBM模型...")
        start_time = datetime.now()

        try:
            # 从配置获取模型参数
            model_params = self.config['model_params']
            cv_params = self.config['data_params']
            training_params = self.config['training_params']

            # 创建LightGBM模型
            self.model = lgb.LGBMRegressor(**model_params)

            # 交叉验证
            if training_params['cross_validation']:
                print(f"执行 {cv_params['cv_splits']} 折交叉验证...")
                kf = KFold(n_splits=cv_params['cv_splits'],
                           shuffle=cv_params['shuffle'],
                           random_state=model_params['random_state'])

                # 手动执行交叉验证以收集详细结果
                fold_results = []
                rmse_folds = []
                mae_folds = []
                r2_folds = []

                for fold, (train_idx, val_idx) in enumerate(kf.split(self.X_train), 1):
                    print(f"训练第 {fold} 折...")

                    X_train_fold = self.X_train.iloc[train_idx]
                    X_val_fold = self.X_train.iloc[val_idx]
                    y_train_fold = self.y_train.iloc[train_idx]
                    y_val_fold = self.y_train.iloc[val_idx]

                    # 训练折模型
                    fold_model = lgb.LGBMRegressor(**model_params)

                    # 如果启用早停，需要验证集
                    if training_params['early_stopping']:
                        fold_model.fit(
                            X_train_fold, y_train_fold,
                            eval_set=[(X_val_fold, y_val_fold)],
                            eval_metric='rmse',
                            callbacks=[lgb.early_stopping(training_params['early_stopping_rounds'], verbose=False)]
                        )
                    else:
                        fold_model.fit(X_train_fold, y_train_fold)

                    # 预测
                    y_pred_fold = fold_model.predict(X_val_fold)

                    # 计算指标
                    rmse = np.sqrt(mean_squared_error(y_val_fold, y_pred_fold))
                    mae = mean_absolute_error(y_val_fold, y_pred_fold)
                    r2 = r2_score(y_val_fold, y_pred_fold)

                    rmse_folds.append(rmse)
                    mae_folds.append(mae)
                    r2_folds.append(r2)

                    # 收集fold结果
                    for i, (true_val, pred_val) in enumerate(zip(y_val_fold, y_pred_fold)):
                        fold_results.append({
                            'fold': fold,
                            'Id': int(val_idx[i]),
                            'y_true': float(true_val),
                            'y_pred': float(pred_val)
                        })

                    print(f"第 {fold} 折完成 - RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}")

                    # 显示最佳迭代次数（如果启用早停）
                    if training_params['early_stopping']:
                        print(f"   最佳迭代次数: {fold_model.best_iteration_}")

                # 记录指标
                self.metrics = {
                    'RMSE': {
                        'mean': float(np.mean(rmse_folds)),
                        'std': float(np.std(rmse_folds)),
                        'folds': [float(score) for score in rmse_folds]
                    },
                    'MAE': {
                        'mean': float(np.mean(mae_folds)),
                        'std': float(np.std(mae_folds)),
                        'folds': [float(score) for score in mae_folds]
                    },
                    'R2': {
                        'mean': float(np.mean(r2_folds)),
                        'std': float(np.std(r2_folds)),
                        'folds': [float(score) for score in r2_folds]
                    }
                }

                # 计算训练时间
                end_time = datetime.now()
                training_time = str(end_time - start_time).split('.')[0]

                print(f"\n交叉验证结果:")
                print(f"RMSE: {self.metrics['RMSE']['mean']:.4f} ± {self.metrics['RMSE']['std']:.4f}")
                print(f"MAE: {self.metrics['MAE']['mean']:.4f} ± {self.metrics['MAE']['std']:.4f}")
                print(f"R2: {self.metrics['R2']['mean']:.4f} ± {self.metrics['R2']['std']:.4f}")
                print(f"训练时间: {training_time}")

                # 生成标准验证报告
                if self.config['output_params']['generate_report']:
                    try:
                        print("生成标准验证报告...")

                        # 临时切换到项目根目录
                        original_cwd = os.getcwd()
                        script_dir = os.path.dirname(os.path.abspath(__file__))
                        model_dir = os.path.dirname(script_dir)
                        project_root = os.path.dirname(model_dir)
                        os.chdir(project_root)

                        generate_validation_outputs(
                            folds_results=fold_results,
                            metrics_dict=self.metrics,
                            model_name=self.config['model_info']['model_name'],
                            model_version=self.config['model_info']['version'],
                            framework=self.config['model_info']['framework'],
                            target=self.config['data_params']['target_variable'],
                            train_size=len(self.X_train),
                            n_features=len(self.feature_names),
                            n_splits=cv_params['cv_splits'],
                            random_state=model_params['random_state'],
                            training_time=training_time,
                            notes="LightGBM model with gradient boosting and early stopping.",
                            args=model_params
                        )

                        # 恢复工作目录
                        os.chdir(original_cwd)

                        # 移动验证报告到outputs/reports目录
                        framework_name = self.config['model_info']['framework']
                        validation_files = [
                            f'validation_{framework_name}.csv',
                            f'validation_report_{framework_name}.json'
                        ]

                        for file_name in validation_files:
                            src_path = os.path.join(project_root, file_name)
                            dst_path = os.path.join(self.output_dirs['reports'], file_name)
                            if os.path.exists(src_path):
                                import shutil
                                shutil.move(src_path, dst_path)
                                print(f"[OK] 移动验证报告到: {dst_path}")

                    except Exception as e:
                        print(f"[WARNING] 标准验证报告生成失败: {e}")

            # 在全部数据上训练最终模型
            print("在全部训练数据上训练最终LightGBM模型...")

            # 如果启用早停，需要分割部分数据作为验证集
            if training_params['early_stopping']:
                from sklearn.model_selection import train_test_split
                X_train_final, X_val_final, y_train_final, y_val_final = train_test_split(
                    self.X_train, self.y_train, test_size=0.1, random_state=model_params['random_state']
                )

                self.model.fit(
                    X_train_final, y_train_final,
                    eval_set=[(X_val_final, y_val_final)],
                    eval_metric='rmse',
                    callbacks=[lgb.early_stopping(training_params['early_stopping_rounds'], verbose=True)]
                )

                print(f"最终模型最佳迭代次数: {self.model.best_iteration_}")
            else:
                self.model.fit(self.X_train, self.y_train)

            # 获取特征重要性
            self.feature_importances = self.model.feature_importances_

            # 打印特征重要性
            feature_importance_df = pd.DataFrame({
                'feature': self.feature_names,
                'importance': self.feature_importances
            }).sort_values('importance', ascending=False)

            print(f"\n[OK] 模型训练完成!")
            print(f"前10个最重要特征:")
            print(feature_importance_df.head(10))

            return True

        except Exception as e:
            print(f"[ERROR] 模型训练失败: {e}")
            import traceback
            traceback.print_exc()
            return False

    def generate_predictions(self):
        """生成预测"""
        print("[INFO] 生成预测...")

        try:
            if self.model is None:
                print("[ERROR] 模型尚未训练")
                return False

            # 预测并还原对数变换
            y_pred_log = self.model.predict(self.X_test)
            y_pred = np.expm1(y_pred_log)

            # 生成提交文件
            submission = pd.DataFrame({
                'Id': self.test_ids,
                'SalePrice': y_pred
            })

            submission_filename = f"submission_{self.config['model_info']['model_name'].lower()}.csv"
            submission_path = os.path.join(os.getcwd(), self.output_dirs['submissions'], submission_filename)
            os.makedirs(os.path.dirname(submission_path), exist_ok=True)
            submission.to_csv(submission_path, index=False)

            print(f"[OK] 提交文件生成成功: {submission_path}")
            print(f"   预测结果统计:")
            print(f"   最小值: ${y_pred.min():,.2f}")
            print(f"   最大值: ${y_pred.max():,.2f}")
            print(f"   平均值: ${y_pred.mean():,.2f}")

            return True

        except Exception as e:
            print(f"[ERROR] 预测生成失败: {e}")
            return False

    def save_model(self):
        """保存模型"""
        print("[INFO] 保存模型文件...")

        try:
            # 保存模型到outputs/models目录
            model_filename = f"{self.config['model_info']['model_name'].lower()}_model.pkl"
            model_path = os.path.join(os.getcwd(), self.output_dirs['models'], model_filename)
            print(f"模型保存路径: {model_path}")

            # 确保目录存在
            os.makedirs(os.path.dirname(model_path), exist_ok=True)
            joblib.dump(self.model, model_path)

            # 保存特征名称到outputs/models目录
            feature_path = os.path.join(os.getcwd(), self.output_dirs['models'], 'feature_names.pkl')
            joblib.dump(self.feature_names, feature_path)

            # 保存特征重要性
            if self.feature_importances is not None:
                importance_df = pd.DataFrame({
                    'feature': self.feature_names,
                    'importance': self.feature_importances
                }).sort_values('importance', ascending=False)

                importance_path = os.path.normpath(os.path.join(self.output_dirs['models'], 'feature_importance.csv'))
                importance_df.to_csv(importance_path, index=False)
                print(f"[OK] 特征重要性保存成功: {importance_path}")

            print(f"[OK] 模型保存成功: {model_path}")
            print(f"[OK] 特征文件保存成功: {feature_path}")
            return True

        except Exception as e:
            print(f"[ERROR] 模型保存失败: {e}")
            return False

    def organize_outputs(self):
        """组织输出文件"""
        # 文件已经直接生成到正确的outputs目录结构中，不需要额外的组织
        print("[INFO] 文件已直接保存到outputs目录结构中")
        return True

    def generate_reports(self):
        """生成报告"""
        print("[INFO] 生成分析报告...")

        try:
            # 验证报告已在训练过程中生成
            if os.path.exists(os.path.join(self.output_dirs['reports'], 'validation_report_lightgbm.json')):
                print("[OK] 验证报告已生成")
            else:
                print("[WARNING] 验证报告未找到")

            # 生成特征重要性图
            if self.feature_importances is not None and self.config['output_params']['save_plots']:
                try:
                    import matplotlib.pyplot as plt
                    import seaborn as sns

                    # 创建特征重要性图
                    feature_importance_df = pd.DataFrame({
                        'feature': self.feature_names,
                        'importance': self.feature_importances
                    }).sort_values('importance', ascending=False)

                    # 选择前20个特征进行可视化
                    top_features = feature_importance_df.head(20)

                    plt.figure(figsize=(12, 8))
                    sns.barplot(data=top_features, x='importance', y='feature')
                    plt.title('LightGBM - Top 20 Feature Importance')
                    plt.xlabel('Importance (Gain)')
                    plt.tight_layout()

                    plot_path = os.path.join(self.output_dirs['reports'], 'feature_importance.png')
                    plt.savefig(plot_path, dpi=300, bbox_inches='tight')
                    plt.close()
                    print(f"[OK] 特征重要性图保存成功: {plot_path}")

                except Exception as e:
                    print(f"[WARNING] 特征重要性图生成失败: {e}")

            return True
        except Exception as e:
            print(f"[ERROR] 报告生成失败: {e}")
            return False

    def run_training(self):
        """运行完整的训练流程"""
        print(f"[INFO] 开始训练模型: {self.config['model_info']['model_name']}")
        print("="*60)

        steps = [
            ("数据加载", self.load_data),
            ("模型训练", self.train_model),
            ("保存模型", self.save_model),
            ("生成预测", self.generate_predictions),
            ("组织输出", self.organize_outputs),
            ("生成报告", self.generate_reports)
        ]

        for step_name, step_func in steps:
            print(f"\n[INFO] {step_name}...")
            if not step_func():
                print(f"[ERROR] {step_name}失败")
                return False

        return True

    def save_config(self):
        """保存配置文件"""
        config_path = 'model_config.json'
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(self.config, f, indent=2, ensure_ascii=False)
        print(f"✅ 配置文件已保存: {config_path}")

def main():
    """主函数"""
    try:
        runner = LightGBMModelRunner()
        success = runner.run_training()

        if success:
            print("\n" + "="*60)
            print("🎉 LightGBM模型训练完成!")
            print("="*60)
            print("📁 输出文件:")
            print(f"   - 模型文件: {runner.config['model_info']['model_name'].lower()}_model.pkl")
            print(f"   - 提交文件: submission_{runner.config['model_info']['model_name'].lower()}.csv")
            print(f"   - 验证报告: validation_report_lightgbm.json")
            print(f"   - 特征重要性: feature_importance.csv")
            print(f"   - 所有文件已整理到 outputs/ 目录")
        else:
            print("\n[ERROR] 模型训练失败!")
            sys.exit(1)

    except Exception as e:
        print(f"\n[ERROR] 运行过程中发生错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()